"""
一个基于 Jina API 的网络搜索和内容抓取工具，

这个工具在 AI-Trader 项目中主要用于获取股票相关的网络信息，为交易决策提供外部数据支持。
"""
import json
import os
import sys
import logging
from datetime import datetime, timedelta
from typing import List, Dict, Any

from fastmcp import FastMCP
from dotenv import load_dotenv
# 获取当前文件所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 获取项目根目录（当前目录的父目录）
root_dir = os.path.dirname(current_dir)
# 将根目录添加到 Python 路径
if root_dir not in sys.path:
    sys.path.append(root_dir)

from prompts.model_prompt import serch_prompt
from qwen_agent.agents import Assistant

load_dotenv()

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

logger = logging.getLogger(__name__)

class AliMCPService:
    """
    调用通义 Web-Search MCP 服务，联网检索
    """
    def __init__(self):
        # 通义 Web-Search MCP 服务
        # 基于通义实验室 Text-Embedding，GTE-reRank，Query 改写，搜索判定等多种检索模型及语义理解，
        # 串接专业搜索工程框架及各类型实时信息检索工具，提供实时互联网全栈信息检索，提升 LLM 回答准确性及时效性。
        self.sse_url = "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/sse"
        self.api_key = os.getenv("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("未提供 API 密钥！请设置 DASHSCOPE_API_KEY 环境变量。")

    def __call__(self, query: str, search_start_date: str, search_end_date: str) -> List[Dict[str, Any]]:
        print(f"Searching for {query}")
        if not search_start_date:
            # 3个月前
            search_start_date = (datetime.now() - timedelta(days=90)).strftime("%Y-%m-%d")
        if not search_end_date:
            search_end_date = datetime.now().strftime("%Y-%m-%d")
        return_content = []
        try :
            # 检查环境变量
            llm_cfg = {'model': 'qwen3-max'}
            system = serch_prompt.format(query=query, search_start_date=search_start_date, search_end_date=search_end_date)

            # 配置MCP工具
            tools = [{
                "mcpServers": {
                    "amap-maps": {
                        "url": self.sse_url,
                        "headers": {
                            "Authorization": f"Bearer {self.api_key}"
                        }
                    }
                }
            }]

            # 创建智能体
            bot = Assistant(
                llm=llm_cfg,
                name='财经信息收集智能体',
                description='财经信息查询',
                system_message=system,
                function_list=tools,
            )

            messages = []
            messages.append({'role': 'user', 'content': query})

            # 执行查询并收集所有响应
            all_responses = []
            for response in bot.run(messages):
                all_responses.append(response)

            # 提取最终的assistant回复内容
            final_content = ""
            if all_responses:
                last_response = all_responses[-1]
                if isinstance(last_response, list):
                    for item in last_response:
                        if isinstance(item, dict) and item.get('role') == 'assistant' and 'content' in item:
                            final_content = item['content']
                elif isinstance(last_response, dict) and 'content' in last_response:
                    final_content = last_response['content']
            # 输出最终结果
            if final_content:
                # 去掉前后```
                final_content = final_content.strip("```json").strip("```")
                final_content = json.loads(final_content)
                return [final_content]
            else:
                return []
        except Exception as e:
            logger.error(str(e))
            return [{
                'error': str(e)
            }]

mcp = FastMCP("Search")

@mcp.tool()
def get_information(query: str, search_end_date: str) -> str:
    """
    使用搜索工具抓取并以结构化方式返回与指定查询相关的主要内容信息。
    查询内容的日期，不能晚于指定日期。

    参数:
        query: 必填，您想要检索的关键信息或搜索词，将在互联网上搜索最匹配的结果。
        search_end_date: 必填，您想要检索的结束日期，格式为 "YYYY-MM-DD"

    返回:
        一个字符串，和查询主题相关的主要内容信息，结构化内容包括：
        - datas: 数据源的创建日期列表
        - content: 和主题相关的内容（前1000个字符）

        如果抓取失败，则返回相应的错误信息。
    """
    try:
        tool = AliMCPService()
        # 3个月
        search_start_date = (datetime.now() - timedelta(days=90)).strftime("%Y-%m-%d")
        if not search_end_date:
            search_end_date = datetime.now().strftime("%Y-%m-%d")
        results = tool(query, search_start_date=search_start_date, search_end_date=search_end_date)
        print("MCP工具查询结果：\n", json.dumps(results, ensure_ascii=False, indent=2))
        # 检查结果是否为空
        if not results:
            return f"⚠️ 搜索查询 '{query}' 未找到结果。可能是网络问题或 API 限制。"

        # 将结果转换为字符串格式
        formatted_results = []
        for result in results:
            if 'error' in result:
                formatted_results.append(f"错误: {result['error']}")
            else:
                formatted_results.append(f"""
                    topic: {query}
                    datas: {result['datas']}
                    content: {result['content'][:1000]}
                    """)
        if not formatted_results:
            return f"⚠️ 搜索查询 '{query}' 返回了空结果。"

        return "\n".join(formatted_results)

    except Exception as e:
        return f"❌ 搜索工具执行失败: {str(e)}"

if __name__ == "__main__":
    # 使用 streamable-http 运行，支持通过环境变量配置主机和端口以避免冲突
    port = int(os.getenv("SEARCH_HTTP_PORT", "8001"))
    mcp.run(transport="streamable-http", port=port)
    #print(get_information(query = "五粮液 000858 2025年10月 股价走势 基本面分析", search_date="2025-10-29"))